// Copyright (C) Kumo inc. and its affiliates.
// Author: Jeff.li lijippy@163.com
// All rights reserved.
// This program is free software: you can redistribute it and/or modify
// it under the terms of the GNU Affero General Public License as published
// by the Free Software Foundation, either version 3 of the License, or
// (at your option) any later version.
//
// This program is distributed in the hope that it will be useful,
// but WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
// GNU Affero General Public License for more details.
//
// You should have received a copy of the GNU Affero General Public License
// along with this program.  If not, see <https://www.gnu.org/licenses/>.
//

#include <nebula/array/builder_base.h>

#include <cstdint>
#include <type_traits>
#include <vector>

#include <nebula/array/array_base.h>
#include <nebula/array/builder_dict.h>
#include <nebula/array/data.h>
#include <nebula/array/util.h>
#include <nebula/core/buffer.h>
#include <nebula/core/builder.h>
#include <nebula/types/scalar.h>
#include <turbo/utility/status.h>
#include <turbo/log/logging.h>
#include <nebula/core/visit_type_inline.h>

namespace nebula {


    turbo::Status ArrayBuilder::check_array_type(const std::shared_ptr<DataType> &expected_type,
                                               const Array &array, const char *message) {
        if (!expected_type->equals(*array.type())) {
            return turbo::failed_precondition_error(message);
        }
        return turbo::OkStatus();
    }

    turbo::Status ArrayBuilder::check_array_type(Type::type expected_type, const Array &array,
                                               const char *message) {
        if (array.type_id() != expected_type) {
            return turbo::failed_precondition_error(message);
        }
        return turbo::OkStatus();
    }

    turbo::Status ArrayBuilder::trim_buffer(const int64_t bytes_filled, ResizableBuffer *buffer) {
        if (buffer) {
            if (bytes_filled < buffer->size()) {
                // Trim buffer
                TURBO_RETURN_NOT_OK(buffer->resize(bytes_filled));
            }
            // zero the padding
            buffer->zero_padding();
        } else {
            // Null buffers are allowed in place of 0-byte buffers
                    DKCHECK_EQ(bytes_filled, 0);
        }
        return turbo::OkStatus();
    }

    turbo::Status ArrayBuilder::append_to_bitmap(bool is_valid) {
        TURBO_RETURN_NOT_OK(Reserve(1));
        unsafe_append_to_bitmap(is_valid);
        return turbo::OkStatus();
    }

    turbo::Status ArrayBuilder::append_to_bitmap(const uint8_t *valid_bytes, int64_t length) {
        TURBO_RETURN_NOT_OK(Reserve(length));
        unsafe_append_to_bitmap(valid_bytes, length);
        return turbo::OkStatus();
    }

    turbo::Status ArrayBuilder::append_to_bitmap(int64_t num_bits, bool value) {
        TURBO_RETURN_NOT_OK(Reserve(num_bits));
        unsafe_append_to_bitmap(num_bits, value);
        return turbo::OkStatus();
    }

    turbo::Status ArrayBuilder::resize(int64_t capacity) {
        TURBO_RETURN_NOT_OK(check_capacity(capacity));
        capacity_ = capacity;
        return null_bitmap_builder_.resize(capacity);
    }

    namespace {

        template<typename ConstIterator>
        struct AppendScalarImpl {
            template<typename T>
            turbo::Status HandleFixedWidth(const T &) {
                auto builder = turbo::checked_cast<typename TypeTraits<T>::BuilderType *>(builder_);
                TURBO_RETURN_NOT_OK(builder->Reserve(n_repeats_ * (scalars_end_ - scalars_begin_)));

                for (int64_t i = 0; i < n_repeats_; i++) {
                    for (auto it = scalars_begin_; it != scalars_end_; ++it) {
                        const auto &scalar = turbo::checked_cast<const typename TypeTraits<T>::ScalarType &>(*it);
                        if (scalar.is_valid) {
                            builder->unsafe_append(scalar.value);
                        } else {
                            builder->unsafe_append_null();
                        }
                    }
                }
                return turbo::OkStatus();
            }

            template<typename T>
            enable_if_t<has_c_type<T>::value, turbo::Status> Visit(const T &t) {
                return HandleFixedWidth(t);
            }

            turbo::Status Visit(const FixedSizeBinaryType &t) { return HandleFixedWidth(t); }

            turbo::Status Visit(const Decimal128Type &t) { return HandleFixedWidth(t); }

            turbo::Status Visit(const Decimal256Type &t) { return HandleFixedWidth(t); }

            template<typename T>
            enable_if_has_string_view<T, turbo::Status> Visit(const T &) {
                int64_t data_size = 0;
                for (auto it = scalars_begin_; it != scalars_end_; ++it) {
                    const auto &scalar = turbo::checked_cast<const typename TypeTraits<T>::ScalarType &>(*it);
                    if (scalar.is_valid) {
                        data_size += scalar.value->size();
                    }
                }

                auto builder = turbo::checked_cast<typename TypeTraits<T>::BuilderType *>(builder_);
                TURBO_RETURN_NOT_OK(builder->Reserve(n_repeats_ * (scalars_end_ - scalars_begin_)));
                TURBO_RETURN_NOT_OK(builder->reserve_data(n_repeats_ * data_size));

                for (int64_t i = 0; i < n_repeats_; i++) {
                    for (auto it = scalars_begin_; it != scalars_end_; ++it) {
                        const auto &scalar = turbo::checked_cast<const typename TypeTraits<T>::ScalarType &>(*it);
                        if (scalar.is_valid) {
                            builder->unsafe_append(std::string_view{*scalar.value});
                        } else {
                            builder->unsafe_append_null();
                        }
                    }
                }
                return turbo::OkStatus();
            }

            template<typename T>
            enable_if_t<is_list_view_type<T>::value || is_list_like_type<T>::value, turbo::Status> Visit(
                    const T &) {
                auto builder = turbo::checked_cast<typename TypeTraits<T>::BuilderType *>(builder_);
                int64_t num_children = 0;
                for (auto it = scalars_begin_; it != scalars_end_; ++it) {
                    if (!it->is_valid) continue;
                    num_children += turbo::checked_cast<const BaseListScalar &>(*it).value->length();
                }
                TURBO_RETURN_NOT_OK(builder->value_builder()->Reserve(num_children * n_repeats_));

                for (int64_t i = 0; i < n_repeats_; i++) {
                    for (auto it = scalars_begin_; it != scalars_end_; ++it) {
                        if (it->is_valid) {
                            const Array &list = *turbo::checked_cast<const BaseListScalar &>(*it).value;
                            if constexpr (T::type_id == Type::MAP || T::type_id == Type::FIXED_SIZE_LIST) {
                                TURBO_RETURN_NOT_OK(builder->append());
                            } else {
                                TURBO_RETURN_NOT_OK(builder->append(/*is_valid=*/true, list.length()));
                            }
                            for (int64_t i = 0; i < list.length(); i++) {
                                TURBO_MOVE_OR_RAISE(auto scalar, list.get_scalar(i));
                                TURBO_RETURN_NOT_OK(builder->value_builder()->append_scalar(*scalar));
                            }
                        } else {
                            TURBO_RETURN_NOT_OK(builder_->append_null());
                        }
                    }
                }
                return turbo::OkStatus();
            }

            turbo::Status Visit(const StructType &type) {
                auto *builder = turbo::checked_cast<StructBuilder *>(builder_);
                auto count = n_repeats_ * (scalars_end_ - scalars_begin_);
                TURBO_RETURN_NOT_OK(builder->Reserve(count));
                for (int field_index = 0; field_index < type.num_fields(); ++field_index) {
                    TURBO_RETURN_NOT_OK(builder->field_builder(field_index)->Reserve(count));
                }
                for (int64_t i = 0; i < n_repeats_; i++) {
                    for (auto it = scalars_begin_; it != scalars_end_; ++it) {
                        const auto &scalar = turbo::checked_cast<const StructScalar &>(*it);
                        for (int field_index = 0; field_index < type.num_fields(); ++field_index) {
                            if (!scalar.is_valid || !scalar.value[field_index]) {
                                TURBO_RETURN_NOT_OK(builder->field_builder(field_index)->append_null());
                            } else {
                                TURBO_RETURN_NOT_OK(builder->field_builder(field_index)
                                                            ->append_scalar(*scalar.value[field_index]));
                            }
                        }
                        TURBO_RETURN_NOT_OK(builder->append(scalar.is_valid));
                    }
                }
                return turbo::OkStatus();
            }

            turbo::Status Visit(const SparseUnionType &type) { return MakeUnionArray(type); }

            turbo::Status Visit(const DenseUnionType &type) { return MakeUnionArray(type); }

            turbo::Status AppendUnionScalar(const DenseUnionType &type, const Scalar &s,
                                            DenseUnionBuilder *builder) {
                const auto &scalar = turbo::checked_cast<const DenseUnionScalar &>(s);
                const auto scalar_field_index = type.child_ids()[scalar.type_code];
                TURBO_RETURN_NOT_OK(builder->append(scalar.type_code));

                for (int field_index = 0; field_index < type.num_fields(); ++field_index) {
                    auto *child_builder = builder->child_builder(field_index).get();
                    if (field_index == scalar_field_index) {
                        if (scalar.is_valid) {
                            TURBO_RETURN_NOT_OK(child_builder->append_scalar(*scalar.value));
                        } else {
                            TURBO_RETURN_NOT_OK(child_builder->append_null());
                        }
                    }
                }
                return turbo::OkStatus();
            }

            turbo::Status AppendUnionScalar(const SparseUnionType &type, const Scalar &s,
                                            SparseUnionBuilder *builder) {
                // For each scalar,
                //  1. append the type code,
                //  2. append the value to the corresponding child,
                //  3. append null to the other children.
                const auto &scalar = turbo::checked_cast<const SparseUnionScalar &>(s);
                TURBO_RETURN_NOT_OK(builder->append(scalar.type_code));

                for (int field_index = 0; field_index < type.num_fields(); ++field_index) {
                    auto *child_builder = builder->child_builder(field_index).get();
                    if (field_index == scalar.child_id) {
                        if (scalar.is_valid) {
                            TURBO_RETURN_NOT_OK(child_builder->append_scalar(*scalar.value[field_index]));
                        } else {
                            TURBO_RETURN_NOT_OK(child_builder->append_null());
                        }
                    } else {
                        TURBO_RETURN_NOT_OK(child_builder->append_null());
                    }
                }
                return turbo::OkStatus();
            }

            template<typename T>
            turbo::Status MakeUnionArray(const T &type) {
                using BuilderType = typename TypeTraits<T>::BuilderType;

                auto *builder = turbo::checked_cast<BuilderType *>(builder_);
                const auto count = n_repeats_ * (scalars_end_ - scalars_begin_);

                TURBO_RETURN_NOT_OK(builder->Reserve(count));

                        DKCHECK_EQ(type.num_fields(), builder->num_children());
                for (int field_index = 0; field_index < type.num_fields(); ++field_index) {
                    TURBO_RETURN_NOT_OK(builder->child_builder(field_index)->Reserve(count));
                }

                for (int64_t i = 0; i < n_repeats_; i++) {
                    for (auto it = scalars_begin_; it != scalars_end_; ++it) {
                        TURBO_RETURN_NOT_OK(AppendUnionScalar(type, *it, builder));
                    }
                }
                return turbo::OkStatus();
            }

            turbo::Status Visit(const RunEndEncodedType &) {
                auto builder = turbo::checked_cast<RunEndEncodedBuilder *>(builder_);

                TURBO_RETURN_NOT_OK(builder->Reserve(n_repeats_ * (scalars_end_ - scalars_begin_)));

                for (int64_t i = 0; i < n_repeats_; i++) {
                    for (auto it = scalars_begin_; it != scalars_end_; ++it) {
                        if (it->is_valid) {
                            const auto &scalar_value = *turbo::checked_cast<const RunEndEncodedScalar &>(*it).value;
                            TURBO_RETURN_NOT_OK(builder->append_scalar(scalar_value, 1));
                        } else {
                            TURBO_RETURN_NOT_OK(builder_->append_null());
                        }
                    }
                }
                return turbo::OkStatus();
            }

            turbo::Status Visit(const DataType &type) {
                return turbo::unimplemented_error("append_scalar for type ", type);
            }

            turbo::Status Convert() { return visit_type_inline(*scalars_begin_->type, this); }

            ConstIterator scalars_begin_;
            ConstIterator scalars_end_;
            int64_t n_repeats_;
            ArrayBuilder *builder_;
        };

// Wraps a const_iterator that has a pointer (or pointer-like) to Scalar as the
// value_type and turns it into an iterator with Scalar as value_type.
        template<typename ConstIterator>
        struct DerefConstIterator {
            ConstIterator it;

            using value_type = Scalar;
            using pointer = const Scalar *;
            using difference_type = typename ConstIterator::difference_type;

            const value_type &operator*() const { return *(*it); }

            DerefConstIterator &operator++() {
                ++it;
                return *this;
            }

            difference_type operator-(const DerefConstIterator &other) const {
                return it - other.it;
            }

            bool operator!=(const DerefConstIterator &other) const { return it != other.it; }

            pointer operator->() const { return &(**it); }
        };

    }  // namespace

    turbo::Status ArrayBuilder::append_scalar(const Scalar &scalar, int64_t n_repeats) {
        if (!scalar.type->equals(type())) {
            return turbo::invalid_argument_error("Cannot append scalar of type ", scalar.type->to_string(),
                                                 " to builder for type ", type()->to_string());
        }
        return AppendScalarImpl<const Scalar *>{&scalar, &scalar + 1, n_repeats, this}.Convert();
    }

    turbo::Status ArrayBuilder::append_scalars(const ScalarVector &scalars) {
        if (scalars.empty()) return turbo::OkStatus();
        const auto ty = type();
        for (const auto &scalar: scalars) {
            if (!scalar->type->equals(ty)) {
                return turbo::invalid_argument_error("Cannot append scalar of type ", scalar->type->to_string(),
                                                     " to builder for type ", type()->to_string());
            }
        }

        using Iterator = DerefConstIterator<ScalarVector::const_iterator>;
        return AppendScalarImpl<Iterator>{Iterator{scalars.begin()}, Iterator{scalars.end()},
                /*n_repeats=*/1, this}
                .Convert();
    }

    turbo::Status ArrayBuilder::finish(std::shared_ptr<Array> *out) {
        std::shared_ptr<ArrayData> internal_data;
        TURBO_RETURN_NOT_OK(finish_internal(&internal_data));
        *out = make_array(internal_data);
        return turbo::OkStatus();
    }

    turbo::Result<std::shared_ptr<Array>> ArrayBuilder::finish() {
        std::shared_ptr<Array> out;
        TURBO_RETURN_NOT_OK(finish(&out));
        return out;
    }

    void ArrayBuilder::reset() {
        capacity_ = length_ = null_count_ = 0;
        null_bitmap_builder_.reset();
    }

    turbo::Status ArrayBuilder::SetNotNull(int64_t length) {
        TURBO_RETURN_NOT_OK(Reserve(length));
        unsafe_set_not_null(length);
        return turbo::OkStatus();
    }

    void ArrayBuilder::unsafe_append_to_bitmap(const std::vector<bool> &is_valid) {
        for (bool element_valid: is_valid) {
            unsafe_append_to_bitmap(element_valid);
        }
    }

    void ArrayBuilder::unsafe_set_not_null(int64_t length) {
        length_ += length;
        null_bitmap_builder_.unsafe_append(length, true);
    }

    void ArrayBuilder::unsafe_set_null(int64_t length) {
        length_ += length;
        null_count_ += length;
        null_bitmap_builder_.unsafe_append(length, false);
    }

}  // namespace nebula
